﻿using System;
using System.Collections.Generic;
using System.IO;
using System.IO.Compression;
using System.Linq;
using System.Reflection;
using dnlib.DotNet;
using dnlib.DotNet.Emit;
using NETDeob.Core.Engine;

namespace NETDeob.Core.Deobfuscators.Confuser._1_7
{
    internal class Resource_Decrypter : IRoutine
    {
        #region Locals

        private readonly List<Int32> _keys = new List<Int32>(3);
        private Resource _encryptedResource;
        private MethodDef _resourceMethod;
        private String _resourceName;

        #endregion

        #region Routine Information

        public string Description
        {
            get { return "Decrypt Resources"; }
        }

        public Dictionary<Func<bool>, string> Phases
        {
            get
            {
                return new Dictionary<Func<bool>, string>
                {
                    {Phase1, "Detection"},
                    {Phase2, "Extraction"},
                    {Phase3, "Decryption"},
                    {Phase4, "Clean Up"}
                };
            }
        }

        #endregion

        #region Phases

        private Boolean Phase1()
        {
            try
            {
                if (!Settings.TargetAssemblyDef.Modules[0].HasResources) return false;

                var typeModule = Settings.TargetAssemblyDef.Modules[0].Types[0]; // <Module>

                if (!typeModule.HasMethods) return false;

                _resourceMethod =
                    typeModule.Methods.FirstOrDefault(
                        m =>
                            m.IsConstructor == false &&
                            m.Body.Instructions.Count >= 1 &&
                            m.ReturnType.FullName == "System.Reflection.Assembly"
                        );

                return _resourceMethod != null;
            }
            catch (Exception exception)
            {
                Logger.Log(exception.Message, 0);
                return false;
            }
        }

        private Boolean Phase2()
        {
            try
            {
                foreach (var instruction in _resourceMethod.Body.Instructions)
                {
                    if (instruction.IsLdcI4())
                        if (instruction.GetLdcI4Value() > 2 && instruction.GetLdcI4Value() < 1000)
                            // TODO: Can be made better
                            _keys.Add(instruction.GetLdcI4Value());


                    if (instruction.OpCode == OpCodes.Ldstr)
                        _resourceName = instruction.Operand.ToString();
                }

                // TODO: Error occurs when Resource Name in Assembly is Shorter than Name in Decryption Method. Temp Fix Applied
                foreach (var resource in Settings.TargetAssemblyDef.Modules[0].Resources)
                    if (resource.Name.ToString().Contains(_resourceName) ||
                        _resourceName.Contains(resource.Name.ToString()))
                        _encryptedResource = resource;


                Logger.Log(String.Concat(
                    "\n", "Resource Keys: ", String.Join(",", _keys.ToArray()),
                    "\n", "Resource Name: ", _resourceName),
                    2);


                return true;
            }
            catch (Exception exception)
            {
                Logger.Exception(exception.Message, 2);
                return false;
            }
        }

        private Boolean Phase3()
        {
            try
            {
                using (var reader =
                        new BinaryReader(
                            new DeflateStream(
                                ((EmbeddedResource) _encryptedResource).GetResourceStream(), CompressionMode.Decompress)))
                {
                    using (var reader2 =
                        new BinaryReader(
                            new DeflateStream(
                                new MemoryStream(
                                    Decrypt(reader.ReadBytes(reader.ReadInt32()))), CompressionMode.Decompress)))
                    {
                        var resourcesAssembly =
                            Assembly.Load(
                                new MemoryStream(
                                    reader2.ReadBytes(reader2.ReadInt32())).ToArray());

                        foreach (var embeddedResourceName in resourcesAssembly.GetManifestResourceNames())
                        {
                            var resourceStream = resourcesAssembly.GetManifestResourceStream(embeddedResourceName);

                            // Prevent Empty Resources from being Added
                            if (resourceStream == null) continue;

                            // Convert Stream to Byte[]
                            var resourceBytes = new Byte[resourceStream.Length];
                            resourceStream.Read(resourceBytes, 0, resourceBytes.Length);

                            // Convert to Embedded Resource
                            Settings.TargetAssemblyDef.Modules[0].Resources.Add(
                                new EmbeddedResource(embeddedResourceName, resourceBytes,
                                    ManifestResourceAttributes.Public));

                            Logger.Log(String.Concat("\n", "Added New Resource:", embeddedResourceName), 2);
                        }
                    }
                }


                return true;
            }
            catch (Exception exception)
            {
                Logger.Exception(exception.Message, 2);
                return false;
            }
        }

        private byte[] Decrypt(IList<Byte> resourceBytes)
        {
            var array = new Byte[resourceBytes.Count/2];
            for (var i = 0; i < resourceBytes.Count; i += 2)
            {
                array[i/2] = (Byte) ((resourceBytes[i + 1] ^ _keys[0])*_keys[1] + (resourceBytes[i] ^ _keys[2]));
            }
            return array;
        }

        private Boolean Phase4()
        {
            try
            {
                // Remove Old Encrypted Resource
                Settings.TargetAssemblyDef.Modules[0].Resources.Remove(_encryptedResource);

                // Remove Decryption method
                Settings.TargetAssemblyDef.Modules[0].Types[0].Methods.Remove(_resourceMethod);

                // Free memory maybe?
                _resourceMethod = null;
                _encryptedResource = null;

                return true;
            }
            catch (Exception exception)
            {
                Logger.Exception(exception.Message, 2);
                return false;
            }
        }

        #endregion
    }
}